package com.ustc.hewei.snowflakemodule.segment;

import com.ustc.hewei.snowflakemodule.exception.OptimisticKeyException;
import com.ustc.hewei.snowflakemodule.segment.entity.ResultId;
import com.ustc.hewei.snowflakemodule.segment.entity.Segment;
import com.ustc.hewei.snowflakemodule.segment.entity.SegmentAlloc;
import com.ustc.hewei.snowflakemodule.segment.entity.SegmentBuffer;
import com.ustc.hewei.snowflakemodule.segment.service.ISegmentAllocService;
import com.ustc.hewei.snowflakemodule.util.Constants;
import com.ustc.hewei.snowflakemodule.util.ThreadPoolFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StopWatch;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author hewei
 * @version 1.0
 * @description: 号段模式id生成
 * @date 2022/11/16 15:26
 */

public class SegmentWorker {
    private static final Logger log = LoggerFactory.getLogger(SegmentWorker.class);
    private final Map<String, SegmentBuffer> bufferCache;
    private final ISegmentAllocService segmentAllocService;
    private final ExecutorService executorService;

    public SegmentWorker(ISegmentAllocService segmentAllocService) {
        this.segmentAllocService = segmentAllocService;
        bufferCache = new ConcurrentHashMap<>();
        executorService = ThreadPoolFactory.createDefaultThreadPool("update_segmentBuffer");
    }

    public boolean init() {
        updateCacheFromDB();
        return true;
    }

    public ResultId getId(String bizTag) {
        SegmentBuffer buffer = bufferCache.get(bizTag);
        if (buffer == null) {
            return ResultId.fail(Constants.UNKNOWN_TAG);
        }
        if (!buffer.getInitRes()) {
            synchronized (buffer) {
                try {
                    if (!buffer.getInitRes()) {
                        updateSegmentFromDB(buffer.getCurrentSegment(), buffer.getBizTag());
                        buffer.setInitRes(true);
                    }
                } catch (OptimisticKeyException e) {
                    throw new RuntimeException(e);
                }
            }
        }
        return getIdByBuffer(buffer);
    }

    public ResultId getIdByBuffer(SegmentBuffer buffer) {
        while (true) {
            buffer.getRLock().lock();
            try {
                Segment segment = buffer.getCurrentSegment();
                if (!buffer.getNextReady() && segment.getRestSpace() < segment.getLoadingId()
                        && buffer.getIsUpdate().compareAndSet(false, true)) {
                    executorService.submit(() -> {
                        try {
                            updateSegmentFromDB(buffer.getNextSegment(), buffer.getBizTag());
                        } catch (OptimisticKeyException e) {
                            throw new RuntimeException(e);
                        } finally {
                            buffer.getWLock().lock();
                            buffer.getIsUpdate().set(false);
                            buffer.setNextReady(true);
                            buffer.getWLock().unlock();
                        }
                    });
                }
                long nextId = segment.getValue().getAndIncrement();
                if (nextId < segment.getMaxId()) {
                    return ResultId.ok(nextId);
                }
            } finally {
                buffer.getRLock().unlock();
            }
            waitForUpdate(buffer);
            buffer.getWLock().lock();
            try {
                Segment currentSegment = buffer.getCurrentSegment();
                long nextId = currentSegment.getValue().getAndIncrement();
                if (nextId < currentSegment.getMaxId()) {
                    return ResultId.ok(nextId);
                }
                if (buffer.getNextReady()) {
                    buffer.switchPos();
                    buffer.setNextReady(false);
                } else {
                    log.error("无可用segment! {}", buffer);
                    //return ResultId.fail(Constants.NO_AVAILABLE_SEGMENT);
                    throw new RuntimeException("无可用segment!");
                }
            } finally {
                buffer.getWLock().unlock();
            }
        }
    }

    public void waitForUpdate(SegmentBuffer buffer) {
        int count = 1;
        while (buffer.getIsUpdate().get()) {
            ++count;
            if (count > 10000) {
                try {
                    TimeUnit.MILLISECONDS.sleep(10);
                    //break;
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    public void updateSegmentFromDB(Segment segment, String bizTag) throws OptimisticKeyException {
        StopWatch stopWatch = new StopWatch();
        stopWatch.start();
        SegmentBuffer buffer = segment.getBuffer();
        SegmentAlloc segmentAlloc;
        if (!buffer.getInitRes()) {
            segmentAlloc = segmentAllocService.updateMaxIdAndGetAlloc(bizTag);
            buffer.setStep(segmentAlloc.getStep());
            buffer.setStep(segmentAlloc.getStep());
        } else if (buffer.getUpdateTime() == 0) {
            segmentAlloc = segmentAllocService.updateMaxIdAndGetAlloc(bizTag);
            buffer.setStep(segmentAlloc.getStep());
            buffer.setStep(segmentAlloc.getStep());
            buffer.setUpdateTime(System.currentTimeMillis());
            segment.setValue(new AtomicLong(0));
        } else {
            long deltaTime = System.currentTimeMillis() - buffer.getUpdateTime();
            int step = buffer.getStep();
            if (deltaTime < Constants.STANDARD_UPDATE_TIME) {
                step = Math.min(step * 2, Constants.MAX_STEP);
            } else if (deltaTime > Constants.STANDARD_UPDATE_TIME) {
                step = Math.max(buffer.getMinStep(), step / 2);
            }
            segmentAlloc = segmentAllocService.updateMaxIdByNewStepWithOptimistic(bizTag, step);
            buffer.setStep(step);
            buffer.setMinStep(segmentAlloc.getStep());
            buffer.setUpdateTime(System.currentTimeMillis());
        }
        segment.getValue().set(segmentAlloc.getMaxId() - buffer.getStep());
        segment.setMaxId(segmentAlloc.getMaxId());
        segment.setStep(buffer.getStep());
        segment.setLoadingId((long) (Constants.UPDATE_RATE * segment.getStep()));
        stopWatch.stop();
        log.info("更新segment {}， 耗时 {}", segment, stopWatch.getLastTaskTimeMillis());
    }

    public void updateCacheFromDB() {
        List<SegmentAlloc> list = segmentAllocService.list();
        for (SegmentAlloc segmentAlloc : list) {
            String bizTag = segmentAlloc.getBizTag();
            SegmentBuffer buffer = new SegmentBuffer();
            buffer.setBizTag(bizTag);
            Segment segment = buffer.getCurrentSegment();
            segment.setMaxId(0L);
            segment.setStep(0);
            segment.setValue(new AtomicLong(0));
            segment.setLoadingId(0L);
            bufferCache.put(bizTag, buffer);
        }
    }
}
